#include "ImageProcess.h"

int SetImageInfo (int W, int H, int SubW, int SubH, ImageInfo *ptrInfo, ImageInfo *ptrSubInfo)
{
    ptrInfo->Width = W;
    ptrInfo->Height = H;
    ptrInfo->Size = W * H;

    ptrSubInfo->Width = SubW;
    ptrSubInfo->Height = SubH;
    ptrSubInfo->Size = SubW * SubH;
    
    return 1;
}

int SetPictureInfo (int W, int H, int SubW, int SubH, PictureInfo *ptrInfo, PictureInfo *ptrSubInfo)
{
    //YUV420
    SetImageInfo (W, H, SubW, SubH, &ptrInfo->Y, &ptrSubInfo->Y);
    SetImageInfo (W/2, H/2, SubW/2, SubH/2, &ptrInfo->U, &ptrSubInfo->U);
    SetImageInfo (W/2, H/2, SubW/2, SubH/2, &ptrInfo->V, &ptrSubInfo->V);

    return 1;
}

int InitImage (ImageInfo strImg, ImageInfo strSubImg, PImage ptrImg, PSubImage ptrSubImg)
{
    int i;
    int numSubImg;

    memcpy (&ptrImg->Info, &strImg, sizeof(ImageInfo));
    ptrImg->ImgBuf = (P_BYTE) malloc (strImg.Size * sizeof(BYTE));

    memcpy (&ptrSubImg->Info, &strSubImg, sizeof(ImageInfo));
    numSubImg = strImg.Size / strSubImg.Size;
    ptrSubImg->NumImg = numSubImg;
    ptrSubImg->ptrImgBuf = (P_BYTE*) malloc (numSubImg * sizeof(P_BYTE*));
    for (i = 0; i < numSubImg; i++)
    {
        ptrSubImg->ptrImgBuf[i] = (P_BYTE) malloc (strSubImg.Size * sizeof(BYTE));
    }
    return 1;
}

int DeleteImage (PImage ptrImg, PSubImage ptrSubImg)
{
    int i;

    if (NULL != ptrImg->ImgBuf)
    {
        free (ptrImg->ImgBuf);
        ptrImg->ImgBuf = NULL;
    }

    for (i = 0; i < ptrSubImg->NumImg; i++)
    {
        if (NULL != ptrSubImg->ptrImgBuf[i])
        {
            free (ptrSubImg->ptrImgBuf[i]);
        }
    }
    free (ptrSubImg->ptrImgBuf);
    ptrSubImg->ptrImgBuf = NULL;

    return 1;
}

int ImageSplit (PImage ptrImg, PSubImage ptrSubImg)
{
    int i, j, hStride, hSubStride;
    int WidthInSubImg;
    P_BYTE ptrSrc, ptrDst, ptr;

    WidthInSubImg = ptrImg->Info.Width / ptrSubImg->Info.Width;
    hStride = ptrImg->Info.Width;
    hSubStride = ptrSubImg->Info.Width;
    ptrSrc = ptrImg->ImgBuf;

    for (i = 0; i < ptrSubImg->NumImg; i++)
    {
        if (i > 0 && i % WidthInSubImg == 0)
        {
            ptrSrc += WidthInSubImg * ptrSubImg->Info.Size - hStride;
        }

        ptr = ptrSrc;
        ptrDst = ptrSubImg->ptrImgBuf[i];
        for (j = 0; j < ptrSubImg->Info.Height; j++)
        {
            memcpy (ptrDst, ptr, hSubStride * sizeof(BYTE));
            ptr += hStride;
            ptrDst += hSubStride;
        }

        ptrSrc += hSubStride;
    }

    return 1;
}

int ImageAssemble (PSubImage ptrSubImg, PImage ptrImg)
{
    int i, j, k, hStride, hSubStride, vSubStride;
    int offset;
    int WidthInSubImg;
    P_BYTE ptrSrc, ptrDst, ptr;
    P_BYTE* ptrBuf;

    WidthInSubImg = ptrImg->Info.Width / ptrSubImg->Info.Width;
    vSubStride = ptrSubImg->Info.Height;
    hStride = ptrImg->Info.Width;
    hSubStride = ptrSubImg->Info.Width;
    
    ptrBuf = ptrSubImg->ptrImgBuf;
    ptrDst = ptrImg->ImgBuf;
    
    for (i = 0; i < ptrSubImg->NumImg; i += WidthInSubImg)
    {
        offset = 0;
        for (j = 0; j < vSubStride; j++)
        {
            for (k = 0; k < WidthInSubImg; k++)
            {
                ptrSrc = ptrBuf[k] + offset;
                memcpy (ptrDst, ptrSrc, hSubStride * sizeof(BYTE));
                ptrDst += hSubStride;
            }
            offset += hSubStride;
        }
        ptrBuf += WidthInSubImg;
    }

    return 1;
}

int InitPicture (PictureInfo strPic, PictureInfo strSubPic, PPicture ptrPic, PSubPicture ptrSubPic)
{
    InitImage (strPic.Y, strSubPic.Y, &ptrPic->Y, &ptrSubPic->Y);
    InitImage (strPic.U, strSubPic.U, &ptrPic->U, &ptrSubPic->U);
    InitImage (strPic.V, strSubPic.V, &ptrPic->V, &ptrSubPic->V);
    return 1;
}

int DeletePicture (PPicture ptrPic, PSubPicture ptrSubPic)
{
    DeleteImage (&ptrPic->Y, &ptrSubPic->Y);
    DeleteImage (&ptrPic->U, &ptrSubPic->U);
    DeleteImage (&ptrPic->V, &ptrSubPic->V);
    return 1;
}

int PictureSplit (PPicture ptrPic, PSubPicture ptrSubPic)
{
    ImageSplit (&ptrPic->Y, &ptrSubPic->Y);
    ImageSplit (&ptrPic->U, &ptrSubPic->U);
    ImageSplit (&ptrPic->V, &ptrSubPic->V);
    return 1;
}

int PictureAssemble (PSubPicture ptrSubPic, PPicture ptrPic)
{
    ImageAssemble (&ptrSubPic->Y, &ptrPic->Y);
    ImageAssemble (&ptrSubPic->U, &ptrPic->U);
    ImageAssemble (&ptrSubPic->V, &ptrPic->V);
    return 1;
}

int LoadPicture (FILE *fp, PPicture ptrPic)
{
    if (NULL != fp)
    {
        fread (ptrPic->Y.ImgBuf, sizeof(BYTE), ptrPic->Y.Info.Size, fp);
        fread (ptrPic->U.ImgBuf, sizeof(BYTE), ptrPic->U.Info.Size, fp);
        fread (ptrPic->V.ImgBuf, sizeof(BYTE), ptrPic->V.Info.Size, fp);
    }
    else
    {
        return 0;
    }
    return 1;
}

int LoadSubPicture (FILE **fp, PSubPicture ptrSubPic)
{
    int i, NumSubImg;

    NumSubImg = ptrSubPic->Y.NumImg;

    for (i = 0; i < NumSubImg; i++)
    {
        if (NULL != fp[i])
        {
            fread (ptrSubPic->Y.ptrImgBuf[i], sizeof(BYTE), ptrSubPic->Y.Info.Size, fp[i]);
            fread (ptrSubPic->U.ptrImgBuf[i], sizeof(BYTE), ptrSubPic->U.Info.Size, fp[i]);
            fread (ptrSubPic->V.ptrImgBuf[i], sizeof(BYTE), ptrSubPic->V.Info.Size, fp[i]);
        }
        else
        {
            return 0;
        }
    }
    return 1;
}

int SavePicture (PPicture ptrPic, FILE *fp)
{
    if (NULL != fp)
    {
        fwrite (ptrPic->Y.ImgBuf, sizeof(BYTE), ptrPic->Y.Info.Size, fp);
        fwrite (ptrPic->U.ImgBuf, sizeof(BYTE), ptrPic->U.Info.Size, fp);
        fwrite (ptrPic->V.ImgBuf, sizeof(BYTE), ptrPic->V.Info.Size, fp);
    }
    else
    {
        return 0;
    }
    return 1;
}

int SaveSubPicture (PSubPicture ptrSubPic, FILE **fp)
{
    int i, NumSubImg;

    NumSubImg = ptrSubPic->Y.NumImg;

    for (i = 0; i < NumSubImg; i++)
    {
        if (NULL != fp[i])
        {
            fwrite (ptrSubPic->Y.ptrImgBuf[i], sizeof(BYTE), ptrSubPic->Y.Info.Size, fp[i]);
            fwrite (ptrSubPic->U.ptrImgBuf[i], sizeof(BYTE), ptrSubPic->U.Info.Size, fp[i]);
            fwrite (ptrSubPic->V.ptrImgBuf[i], sizeof(BYTE), ptrSubPic->V.Info.Size, fp[i]);
        }
        else
        {
            return 0;
        }
    }
    return 1;
}

int Split (int argc, char** argv)
{
    PictureInfo Info;
    PictureInfo SubInfo;
    Picture Pic;
    SubPicture SubPic;

    FILE *fpIn, **fpOut;
    char SubFileName[500];
    int i, FrmNum;

    FrmNum = atoi(argv[0]);

    SetPictureInfo (atoi(argv[1]), atoi(argv[2]), atoi(argv[3]), atoi(argv[4]), &Info, &SubInfo);
    
    InitPicture (Info, SubInfo, &Pic, &SubPic);

    fpOut = (FILE **) malloc(SubPic.Y.NumImg * sizeof(FILE*));

    fpIn = fopen(argv[5], "rb");

    for (i = 0; i < SubPic.Y.NumImg; i++)
    {
        sprintf (SubFileName, "%s_%d.yuv", argv[6], i);
        fpOut[i] = fopen (SubFileName, "wb");
    }
    
    for (i = 0; i < FrmNum; i++)
    {
        LoadPicture (fpIn, &Pic);
        PictureSplit (&Pic, &SubPic);
        SaveSubPicture (&SubPic, fpOut);
    }

    fclose (fpIn);
    for (i = 0; i < SubPic.Y.NumImg; i++)
    {
        fclose (fpOut[i]);
    }
    free(fpOut);

    DeletePicture (&Pic, &SubPic);
    
    return 1;
}

int Assemble (int argc, char** argv)
{
    PictureInfo Info;
    PictureInfo SubInfo;
    Picture Pic;
    SubPicture SubPic;

    FILE **fpIn, *fpOut;
    char SubFileName[500];
    int i, FrmNum;

    FrmNum = atoi(argv[0]);

    SetPictureInfo (atoi(argv[1]), atoi(argv[2]), atoi(argv[3]), atoi(argv[4]), &Info, &SubInfo);
    
    InitPicture (Info, SubInfo, &Pic, &SubPic);

    fpIn = (FILE**) malloc(SubPic.Y.NumImg * sizeof(FILE*));

    for (i = 0; i < SubPic.Y.NumImg; i++)
    {
        sprintf (SubFileName, "%s_%d.yuv", argv[5], i);
        fpIn[i] = fopen (SubFileName, "rb");
    }

    fpOut = fopen(argv[6], "wb");
    
    for (i = 0; i < FrmNum; i++)
    {
        LoadSubPicture (fpIn, &SubPic);
        PictureAssemble (&SubPic, &Pic);
        SavePicture (&Pic, fpOut);
    }
    
    for (i = 0; i < SubPic.Y.NumImg; i++)
    {
        fclose (fpIn[i]);
    }
    fclose (fpOut);
    free (fpIn);

    DeletePicture (&Pic, &SubPic);
    
    return 1;
}